use std::time::Duration;

use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        Query, State,
    },
    response::IntoResponse,
};
use futures_util::{SinkExt, StreamExt};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};

use crate::{
    config::validate_token,
    monitor::{
        collect_system_stats, dns_lookup, http_bench, parse_ports, ping, portscan, tcping,
        ws_bench, DnsOptions, HttpBenchOptions, PingOptions, PortscanOptions, TcpingOptions,
        WsBenchOptions,
    },
};

use super::{AppState, AuthQuery};

pub async fn ws_upgrade(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
    Query(query): Query<AuthQuery>,
) -> impl IntoResponse {
    ws.on_upgrade(move |socket| ws_monitor_handler(socket, state, query))
}

pub async fn ws_monitor_handler(socket: WebSocket, state: AppState, query: AuthQuery) {
    let token = query.token.unwrap_or_default();
    let (interval_ms, authorized) = {
        let config = state.config.read().await;
        (
            config.server.interval_ms.max(1),
            validate_token(&config, &token),
        )
    };
    let (mut sender, mut receiver) = socket.split();

    if !authorized {
        let _ = sender
            .send(Message::Text(
                json!({
                    "code": 401,
                    "message": "invalid token",
                    "data": null
                })
                .to_string(),
            ))
            .await;
        let _ = sender.close().await;
        return;
    }

    let mut interval = tokio::time::interval(Duration::from_millis(interval_ms));
    loop {
        tokio::select! {
            _ = interval.tick() => {
                let system_stats = collect_system_stats();
                let network_stats = state.monitor.collect_stats();

                let payload = json!({
                    "code": 200,
                    "message": "success",
                    "data": {
                        "system": system_stats,
                        "network": network_stats,
                    },
                    "timestamp": chrono::Utc::now().timestamp_millis(),
                });

                if sender.send(Message::Text(payload.to_string())).await.is_err() {
                    break;
                }
            }
            msg = receiver.next() => {
                match msg {
                    Some(Ok(Message::Close(frame))) => {
                        let _ = sender.send(Message::Close(frame)).await;
                        break;
                    }
                    Some(Ok(Message::Ping(payload))) => {
                        if sender.send(Message::Pong(payload)).await.is_err() {
                            break;
                        }
                    }
                    Some(Ok(Message::Text(_))) | Some(Ok(Message::Binary(_))) | Some(Ok(Message::Pong(_))) => {}
                    Some(Err(_)) | None => break,
                }
            }
        }
    }

    let _ = sender.close().await;
}

#[derive(Deserialize)]
#[serde(tag = "action", rename_all = "lowercase")]
enum ProbeCommand {
    Ping {
        target: String,
        #[serde(default = "default_ping_count")]
        count: u32,
        #[serde(default = "default_ping_timeout")]
        timeout_ms: u64,
    },
    Tcping {
        target: String,
        port: u16,
        #[serde(default = "default_tcping_count")]
        count: u32,
        #[serde(default = "default_tcping_timeout")]
        timeout_ms: u64,
    },
    Dns {
        domain: String,
        #[serde(default = "default_dns_record_type")]
        record_type: String,
    },
    Portscan {
        target: String,
        ports: String,
        #[serde(default = "default_portscan_timeout")]
        timeout_ms: u64,
        #[serde(default = "default_portscan_concurrency")]
        concurrency: usize,
    },
    HttpBench {
        url: String,
        #[serde(default = "default_http_requests")]
        requests: u32,
        #[serde(default = "default_http_concurrency")]
        concurrency: usize,
        #[serde(default = "default_http_timeout")]
        timeout_ms: u64,
    },
    WsBench {
        url: String,
        #[serde(default = "default_ws_connections")]
        connections: u32,
        #[serde(default = "default_ws_concurrency")]
        concurrency: usize,
        #[serde(default = "default_ws_messages")]
        messages: u32,
        payload: Option<String>,
        #[serde(default = "default_ws_timeout")]
        timeout_ms: u64,
    },
}

fn default_ping_count() -> u32 {
    4
}

fn default_ping_timeout() -> u64 {
    1000
}

fn default_tcping_count() -> u32 {
    4
}

fn default_tcping_timeout() -> u64 {
    1000
}

fn default_dns_record_type() -> String {
    "A".into()
}

fn default_portscan_timeout() -> u64 {
    500
}

fn default_portscan_concurrency() -> usize {
    50
}

fn default_http_requests() -> u32 {
    100
}

fn default_http_concurrency() -> usize {
    10
}

fn default_http_timeout() -> u64 {
    1000
}

fn default_ws_connections() -> u32 {
    50
}

fn default_ws_concurrency() -> usize {
    10
}

fn default_ws_messages() -> u32 {
    0
}

fn default_ws_timeout() -> u64 {
    1000
}

#[derive(Serialize)]
struct ProbeResponse {
    code: u16,
    message: String,
    action: String,
    data: Option<Value>,
    error: Option<String>,
    elapsed_ms: Option<u128>,
    timestamp: i64,
}

pub async fn ws_probe_upgrade(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
    Query(query): Query<AuthQuery>,
) -> impl IntoResponse {
    ws.on_upgrade(move |socket| ws_probe_handler(socket, state, query))
}

pub async fn ws_probe_handler(socket: WebSocket, state: AppState, query: AuthQuery) {
    let token = query.token.unwrap_or_default();
    let (mut sender, mut receiver) = socket.split();

    let authorized = {
        let config = state.config.read().await;
        validate_token(&config, &token)
    };

    if !authorized {
        let _ = sender
            .send(Message::Text(
                json!({
                    "code": 401,
                    "message": "invalid token",
                    "data": null
                })
                .to_string(),
            ))
            .await;
        let _ = sender.close().await;
        return;
    }

    while let Some(msg) = receiver.next().await {
        match msg {
            Ok(Message::Text(text)) => {
                let cmd = serde_json::from_str::<ProbeCommand>(&text);
                match cmd {
                    Ok(command) => {
                        let action = command_action(&command).to_string();
                        let result = handle_probe_command(command).await;
                        let payload = ProbeResponse {
                            code: if result.error.is_some() { 400 } else { 200 },
                            message: result.error.clone().unwrap_or_else(|| "success".into()),
                            action,
                            data: result.data,
                            error: result.error,
                            elapsed_ms: result.elapsed_ms,
                            timestamp: chrono::Utc::now().timestamp_millis(),
                        };
                        if sender
                            .send(Message::Text(serde_json::to_string(&payload).unwrap()))
                            .await
                            .is_err()
                        {
                            break;
                        }
                    }
                    Err(err) => {
                        let payload = json!({
                            "code": 400,
                            "message": "invalid payload",
                            "error": err.to_string(),
                            "timestamp": chrono::Utc::now().timestamp_millis(),
                        });
                        if sender
                            .send(Message::Text(payload.to_string()))
                            .await
                            .is_err()
                        {
                            break;
                        }
                    }
                }
            }
            Ok(Message::Close(frame)) => {
                let _ = sender.send(Message::Close(frame)).await;
                break;
            }
            Ok(Message::Ping(payload)) => {
                if sender.send(Message::Pong(payload)).await.is_err() {
                    break;
                }
            }
            Ok(Message::Binary(_)) | Ok(Message::Pong(_)) => {}
            Err(_) => break,
        }
    }

    let _ = sender.close().await;
}

fn command_action(command: &ProbeCommand) -> &'static str {
    match command {
        ProbeCommand::Ping { .. } => "ping",
        ProbeCommand::Tcping { .. } => "tcping",
        ProbeCommand::Dns { .. } => "dns",
        ProbeCommand::Portscan { .. } => "portscan",
        ProbeCommand::HttpBench { .. } => "http_bench",
        ProbeCommand::WsBench { .. } => "ws_bench",
    }
}

struct ProbeResult {
    data: Option<Value>,
    error: Option<String>,
    elapsed_ms: Option<u128>,
}

async fn handle_probe_command(command: ProbeCommand) -> ProbeResult {
    match command {
        ProbeCommand::Ping {
            target,
            count,
            timeout_ms,
        } => {
            let envelope = ping(PingOptions {
                target: &target,
                count,
                timeout: Duration::from_millis(timeout_ms),
            })
            .await;
            ProbeResult {
                elapsed_ms: Some(envelope.elapsed_ms),
                error: envelope.error.clone(),
                data: Some(serde_json::to_value(envelope).unwrap()),
            }
        }
        ProbeCommand::Tcping {
            target,
            port,
            count,
            timeout_ms,
        } => {
            let envelope = tcping(TcpingOptions {
                target: &target,
                port,
                count,
                timeout: Duration::from_millis(timeout_ms),
            })
            .await;
            ProbeResult {
                elapsed_ms: Some(envelope.elapsed_ms),
                error: envelope.error.clone(),
                data: Some(serde_json::to_value(envelope).unwrap()),
            }
        }
        ProbeCommand::Dns {
            domain,
            record_type,
        } => {
            let envelope = dns_lookup(DnsOptions {
                query: &domain,
                record_type: &record_type,
            })
            .await;
            ProbeResult {
                elapsed_ms: Some(envelope.elapsed_ms),
                error: envelope.error.clone(),
                data: Some(serde_json::to_value(envelope).unwrap()),
            }
        }
        ProbeCommand::Portscan {
            target,
            ports,
            timeout_ms,
            concurrency,
        } => match parse_ports(&ports) {
            Ok(parsed) => {
                let envelope = portscan(PortscanOptions {
                    target: &target,
                    ports: parsed,
                    timeout: Duration::from_millis(timeout_ms),
                    concurrency: concurrency.max(1),
                })
                .await;
                ProbeResult {
                    elapsed_ms: Some(envelope.elapsed_ms),
                    error: envelope.error.clone(),
                    data: Some(serde_json::to_value(envelope).unwrap()),
                }
            }
            Err(err) => ProbeResult {
                data: None,
                error: Some(err),
                elapsed_ms: None,
            },
        },
        ProbeCommand::HttpBench {
            url,
            requests,
            concurrency,
            timeout_ms,
        } => {
            let envelope = http_bench(HttpBenchOptions {
                url: &url,
                requests,
                concurrency,
                timeout_ms,
            })
            .await;
            ProbeResult {
                elapsed_ms: Some(envelope.elapsed_ms),
                error: envelope.error.clone(),
                data: Some(serde_json::to_value(envelope).unwrap()),
            }
        }
        ProbeCommand::WsBench {
            url,
            connections,
            concurrency,
            messages,
            payload,
            timeout_ms,
        } => {
            let envelope = ws_bench(WsBenchOptions {
                url: &url,
                connections,
                concurrency,
                messages,
                payload: payload.as_deref(),
                timeout_ms,
            })
            .await;
            ProbeResult {
                elapsed_ms: Some(envelope.elapsed_ms),
                error: envelope.error.clone(),
                data: Some(serde_json::to_value(envelope).unwrap()),
            }
        }
    }
}
